rm(list = ls())

library(tidyverse)
library(lfe)
library(readxl)
library(pbapply)
library(broom)

## Helper functions

source("code/helper_functions.R")

## Load Data

btw <- readRDS("data/data_main.RDS") %>%
  dplyr::select(
    county_id_2018,
    year,
    matches("zweit|diff"),
    turnout,
    east,
    treat_categorical4,
    matches("covars")
  ) %>%
  mutate(treat_categorical4 = factor(treat_categorical4,
    levels = c(
      "NoChange",
      "DecAndInc",
      "Increase",
      "Decrease"
    )
  )) %>%
  filter(!year == 2013)

## Def outcomes 

outcomes <- c("small_party_zweit_diff", 
              "turnout_diff",
              "polar_bund_alt_diff", 
              "polar_land_alt_diff",
              "polar_unit_alt_diff")

## Def covariates 

covars <- c("covar_pop_total_diff", "covar_gdp_pc_diff")

## Def subsets

ss_list <- list(
  btw,
  btw %>% filter(year < 2000),
  btw %>% filter(year > 2000),
  btw %>% filter(!(east == 1 & year == 1994))
)

ss_labels <- c(
  "Combined",
  "Pre 2005", "Post 2005",
  "NoEastIn94"
)

##

res <- lapply(outcomes, function(o) {
  lapply(1:length(ss_list), function(i) {
    ## Get DF

    df_temp <- ss_list[[i]]
    ss_label = ss_labels[i]

    ## Prep

    df_temp$outcome <- df_temp %>% pull(!!o)

    df_temp <- df_temp %>%
      filter(!is.na(outcome) & !is.na(treat_categorical4)) %>%
      dplyr::select(
        outcome, year, county_id_2018,
        east,
        treat_categorical4, one_of(covars)
      )


    ## Estimate

    m1 <- felm(outcome ~ treat_categorical4 | year | 0 | county_id_2018,
      data = df_temp
    ) %>%
      tidy_felm() %>%
      mutate(fe = "year")

    ## Year + covars

    f <- paste0(
      "outcome ~ treat_categorical4 +",
      paste0(covars, collapse = "+"),
      "| year | 0 | county_id_2018"
    ) %>%
      as.formula()

    m2 <- felm(f, data = df_temp) %>%
      tidy_felm() %>%
      mutate(fe = "year_covars")

    rbind(m1, m2) %>%
      mutate(
        outcome = !!o,
        ss = ss_label
      )
  }) %>% reduce(rbind)
}) %>%
  reduce(rbind) %>%
  mutate(period = "now") %>%
  filter(!str_detect(term, "Intercept|covar"))

## Prep

results_btw_main <- res %>%
  mutate(term = str_remove(term, "treat_categorical4")) %>%
  filter(!str_detect(term, "Intercept|covar"))

## Lag treatment

btw <- btw %>%
  group_by(county_id_2018) %>%
  arrange(year) %>%
  mutate(
    treat_categorical4_lag1 = lag(treat_categorical4),
    treat_categorical4_lag2 = lag(treat_categorical4, 2),
  ) %>%
  ungroup()

## Lead treatment

btw <- btw %>%
  group_by(county_id_2018) %>%
  arrange(year) %>%
  mutate(
    treat_categorical4_lead1 = lead(treat_categorical4),
    treat_categorical4_lead2 = lead(treat_categorical4, 2),
  ) %>%
  ungroup()

## Number of lags and leads

n_lag <- 2
n_lead <- 2

## Estimate across lags

res <- pblapply(outcomes, function(o) {
    lapply(1:n_lag, function(l) {

        ## Get DF

        df_temp <- btw

        ## Declare treatment

        df_temp$treatvar <- df_temp[, paste0("treat_categorical4_lag", l)] %>%
          pull(1)


        ## Prep

        df_temp$outcome <- df_temp %>%
          pull(!!o)

        df_temp <- df_temp %>%
          filter(!is.na(outcome) & !is.na(treatvar)) %>%
          dplyr::select(
            outcome, year, county_id_2018,
            treatvar,
            one_of(covars)
          )
        
        ## Estimate


        m1 <- felm(outcome ~ treatvar | year | 0 | county_id_2018,
          data = df_temp
        ) %>%
          tidy_felm() %>%
          mutate(fe = "year")

        ## Year + covars

        f <- paste0(
          "outcome ~ treatvar +",
          paste0(covars, collapse = "+"),
          "| year | 0 | county_id_2018"
        ) %>%
          as.formula()

        m2 <- felm(f, data = df_temp) %>%
          tidy_felm() %>%
          mutate(fe = "year_covars")

        ## Combine

        rbind(m1, m2) %>%
          mutate(
            outcome = !!o,
            ss = "Combined"
          ) %>%
          mutate(period = paste0("lag", l))
        
    }) %>% reduce(rbind)
}) %>% reduce(rbind)

## Prep

results_lag <- res %>%
  mutate(term = str_remove(term, "treatvar")) %>%
  filter(!str_detect(term, "Intercept|covar"))

## Do the same for the lead treatment 

res <- pblapply(outcomes, function(o) {
    lapply(1:n_lead, function(l) {

        ## Get DF

        df_temp <- btw

        ## Declare treatment

        df_temp$treatvar <- df_temp[, paste0("treat_categorical4_lead", l)] %>%
          pull(1)
        
        ## Prep

        df_temp$outcome <- df_temp %>%
          pull(!!o)
        
        df_temp <- df_temp %>%
          filter(!is.na(outcome) & !is.na(treatvar)) %>%
          dplyr::select(
            outcome, year, county_id_2018,
            treatvar,
            one_of(covars)
          )
        
        ## Estimate

        m1 <- felm(outcome ~ treatvar | year | 0 | county_id_2018,
          data = df_temp
        ) %>%
          tidy_felm() %>%
          mutate(fe = "year")
        
        ## Year + covars

        f <- paste0(
          "outcome ~ treatvar +",
          paste0(covars, collapse = "+"),
          "| year | 0 | county_id_2018"
        ) %>%
          as.formula()
        
        m2 <- felm(f, data = df_temp) %>%
          tidy_felm() %>%
          mutate(fe = "year_covars")
        
        ## Combine

        rbind(m1, m2) %>%
          mutate(
            east = c("All"),
            outcome = !!o,
            ss = "Combined"
          ) %>%
          mutate(period = paste0("lead", l))
  
    }) %>% reduce(rbind)
}) %>% reduce(rbind)

## Prep

results_lead <- res %>%
  mutate(term = str_remove(term, "treatvar")) %>%
  filter(!str_detect(term, "Intercept|covar"))

## Remove all objects other than results_*

rm(list = ls()[!str_detect(ls(), "results_")])

source("code/helper_functions.R")

## Municipal elections 

muni <- readRDS("data/data_main_municipal.RDS") %>%
  dplyr::select(
    county_id_2018,
    year,
    matches("diff"),
    treat_categorical4,
    matches("covars"),
  ) %>%
  mutate(treat_categorical4 = factor(treat_categorical4,
    levels = c(
      "NoChange",
      "DecAndInc",
      "Increase",
      "Decrease"
    )
  ))

## Def outcomes

outcomes <- c(
  "small_party_diff",
  "turnout_diff",
  "polar_bund_alt_diff",
  "polar_land_alt_diff",
  "polar_unit_alt_diff"
)

## Def covariates

covars <- c("covar_pop_total_diff", "covar_gdp_pc_diff")

##

res <- pblapply(outcomes, function(o) {
  ## Get DF

  df_temp <- muni

  ## Prep

  df_temp$outcome <- df_temp %>%
    pull(!!o)

  df_temp <- df_temp %>%
    filter(!is.na(outcome) & !is.na(treat_categorical4)) %>%
    dplyr::select(
      outcome, year, county_id_2018,
      treat_categorical4,
      one_of(covars)
    )


  ## Estimate

  m1 <- felm(outcome ~ treat_categorical4 | year | 0 | county_id_2018,
    data = df_temp
  ) %>%
    tidy_felm() %>%
    mutate(fe = "year")

  ## Year + covars

  f <- paste0(
    "outcome ~ treat_categorical4 +",
    paste0(covars, collapse = "+"),
    "| year | 0 | county_id_2018"
  ) %>%
    as.formula()

  m2 <- felm(f, data = df_temp) %>%
    tidy_felm() %>%
    mutate(fe = "year_covars")

  ## Combine

  rbind(m1, m2) %>%
    mutate(
      outcome = !!o,
      ss = "Combined"
    )
}) %>%
  reduce(rbind) %>%
  mutate(period = "now") %>%
  filter(!str_detect(term, "Intercept|covar"))

results_muni <- res %>%
  mutate(term = str_remove(term, "treat_categorical4"))

## Combine all federal election results

res_bt <- bind_rows(results_btw_main, results_lag, results_lead) %>%
  mutate(outcome = str_remove(outcome, "_zweit")) %>%
  mutate(unit = "Federal election")

## Municipal elections

results_muni <- results_muni %>%
  mutate(outcome = str_remove(outcome, "_parties")) %>%
  mutate(unit = "Municipal election")

## Combine all

res_full <- bind_rows(res_bt, results_muni) %>%
  mutate(fe = dplyr::recode(fe,
    `year` = "Year FE",
    `year_covars` = "Covars + Year FE"
  )) %>%
  mutate(fe = factor(fe, levels = c(
    "Year FE",
    "Covars + Year FE"
  ))) %>%
  mutate(unit = factor(unit, levels = c(
    "Federal election",
    "Municipal election"
  ))) %>%
  mutate(period_numeric = dplyr::recode(period,
    `lag1` = 2,
    `lag2` = 3,
    `now` = 1,
    `lead1` = -1,
    `lead2` = -2
  )) %>%
  mutate(
    conf.low90 = estimate - qnorm(0.95) * std.error,
    conf.high90 = estimate + qnorm(0.95) * std.error
  )

## Rename the outcome

res_full <- res_full %>%
  filter(!outcome %in% c(
    "polar_bund_diff",
    "polar_land_diff",
    "polar_unit_diff"
  )) %>%
  mutate(outcome = dplyr::recode(outcome,
    `polar_bund_alt_diff` = "polar_bund_diff",
    `polar_land_alt_diff` = "polar_land_diff",
    `polar_unit_alt_diff` = "polar_unit_diff"
  ))

## Def main specification

res_full <- res_full %>%
  mutate(main_spec = case_when(
    fe == "Year FE" ~ 1,
    TRUE ~ 0
  ))

## Save this

write_rds(res_full, "data/Results_Main.rds")

